import GAN
import numpy as np
import matplotlib.pyplot as plt
import matplotlib.image as mimg
import os
import time


class FACE_DCGAN(object):

    def __init__(self):
        self.img_rows = 96
        self.img_cols = 96
        self.channel = 3
        self.dir = 'face/'
        listdir = os.listdir(self.dir)
        self.x_train = []
        t = time.time()
        for i in listdir:
            self.x_train.append(mimg.imread(self.dir + i))
        print(time.time() - t)
        self.x_train = (np.array(self.x_train) / 256).astype('float64')
        self.DCGAN = GAN.DCGAN(96, 96, 3)
        self.discriminator = self.DCGAN.discriminator_model()
        self.adversarial = self.DCGAN.adversarial_model()
        self.generator = self.DCGAN.generator()

    def train(self, train_steps=2000, batch_size=256):
        for i in range(train_steps):
            t = time.time()
            images_train = self.x_train[np.random.randint(
                0, self.x_train.shape[0], size=batch_size), :, :, :]
            noise = np.random.uniform(-1.0, 1.0, size=[batch_size, 100])
            images_fake = self.generator.predict(noise)
            x = np.concatenate((images_train, images_fake))
            y = np.ones([2 * batch_size, 1])
            y[batch_size:, :] = 0
            for j in range(5):
                d_loss = self.discriminator.train_on_batch(x, y)
            y = np.ones([batch_size, 1])
            noise = np.random.uniform(-1.0, 1.0, size=[batch_size, 100])
            a_loss = self.adversarial.train_on_batch(noise, y)
            if i % 50 == 0:
                self.plot(i)
            t = time.time() - t
#            self.plot(6)
            print(i, t)
        self.plot(i)

    def plot(self, index):
        noise = np.random.uniform(-1.0, 1.0, size=[1, 100])
        image = self.generator.predict(noise)
        image = np.reshape(image, [self.img_rows, self.img_cols, self.channel])
        image = (256 * image).astype('int8')
        plt.imshow(image)
        plt.axis('off')
        plt.tight_layout()
        filename = 'face_%d.png' % index
        plt.savefig(filename)
        plt.close('all')
